import numpy as np
import pandas as pd
from numpy import matrix as m
from pandas import DataFrame as df
from scipy import linalg
try:
    import seaborn
except:
    pass
from matplotlib.pylab import rcParams
import matplotlib.pyplot as plt
rcParams['figure.figsize'] = 11, 4

show_flag = False;
class mySSA_to_fit(object):
    '''Singular Spectrum Analysis object'''
    def __init__(self, time_series,steps_ahead):
        # 数据帧(DataFrame)是二维数据结构
        self.org = pd.DataFrame(time_series)
        print(self.org.shape[0])
        #self.ts = pd.DataFrame(time_series[:time_series.shape[0] - steps_ahead])
        self.ts = pd.DataFrame(time_series[:time_series.shape[0]])
        
        self.ts_name = self.ts.columns.tolist()[0]
        if self.ts_name==0:
            self.ts_name = 'ts'
        self.ts_v = self.ts.values
         # 获取行数
        self.ts_N = self.ts.shape[0]
        #MS频率
        self.freq = self.ts.index.inferred_freq
    
        
    @staticmethod
    def _printer(name, *args):
        '''Helper function to print messages neatly'''
        print('-'*40)
        print(name+':')
        for msg in args:
            print(msg)  
    
    @staticmethod
    def _dot(x,y):
        '''Alternative formulation of dot product to allow missing values in arrays/matrices'''
        pass
    
    @staticmethod
    def get_contributions(X=None, s=None, plot=True):
        '''Calculate the relative contribution of each of the singular values'''
        lambdas = np.power(s,2)
        # 根据ord参数的值，该函数能够返回八种不同矩阵范数中的一种，或者无数种矢量范数之一
        # 得到的结果是个值
        # ║A║2 = A的最大奇异值 = (max{ λi(AH*A) }) 1/2 

        # （谱范数，即A^H*A特征值λi中最大者λ1的平方根，其中AH为A的转置共轭矩阵）；
        frob_norm = np.linalg.norm(X)
        
        #贡献值累加
        sum =0
        for element in lambdas.flat:
            sum += element
         #ret = df(lambdas/(frob_norm**2), columns=['Contribution'])
         #frob_norm**2 == sum
        ret = df(lambdas/(sum), columns=['Contribution'])
       # 保留4位小数
        ret['Contribution'] = ret.Contribution.round(4)
        if plot:
            ax = ret[ret.Contribution!=0].plot.bar(legend=False)
            ax.set_xlabel("Lambda_i")
            ax.set_title('Non-zero contributions of Lambda_i')
            vals = ax.get_yticks()
             # 设置y轴格式
            ax.set_yticklabels(['{:5.2f}%'.format(x*100) for x in vals])
            if show_flag:
               plt.show(ax)
            return ax
        return ret[ret.Contribution>0]
    
    @staticmethod
    def diagonal_averaging(hankel_matrix):
        '''Performs anti-diagonal averaging from given hankel matrix
        Returns: Pandas DataFrame object containing the reconstructed series'''
        mat = m(hankel_matrix)
        L, K = mat.shape
        L_star, K_star = min(L,K), max(L,K)
        new = np.zeros((L,K))
        if L > K:
            mat = mat.T
        ret = []
        
        #Diagonal Averaging
        for k in range(1-K_star, L_star):
            mask = np.eye(K_star, k=k, dtype='bool')[::-1][:L_star,:]
            mask_n = sum(sum(mask))
            ma = np.ma.masked_array(mat.A, mask=1-mask)
            ret+=[ma.sum()/mask_n]
        
        return df(ret).rename(columns={0:'Reconstruction'})
        
    def view_time_series(self):
        '''Plot the time series'''
        self.ts.plot(title='Original Time Series')
        if show_flag:
            plt.show()

    # 嵌入
    def embed(self, embedding_dimension=None, suspected_frequency=None, verbose=False, return_df=False):
        '''Embed the time series with embedding_dimension window size.
        Optional: suspected_frequency changes embedding_dimension such that it is divisible by suspected frequency'''
        if not embedding_dimension:
            # 如果没有窗口长度，则取1/2
            self.embedding_dimension = self.ts_N//2
        else:
            self.embedding_dimension = embedding_dimension
        if suspected_frequency:
            self.suspected_frequency = suspected_frequency
            # 取周期的整数倍
            self.embedding_dimension = (self.embedding_dimension//self.suspected_frequency)*self.suspected_frequency
    
        self.K = self.ts_N-self.embedding_dimension+1
        # 构建hankle矩阵，每一条副对角线上的元素相等
        # hankle矩阵如下
        # x0 x1 x2 x3 ..  xk
        # x1 x2 x3    ..
        # x3 x4       ..
        # xL xL+1     ..  xN
        # 即self.X = m(linalg.hankel(self.ts, np.zeros(self.embedding_dimension)))实现的效果
        # T表示转置
        # T[:,:self.K] 表示取所有行的 0 ~ K 列的数据
        self.X = m(linalg.hankel(self.ts, np.zeros(self.embedding_dimension))).T[:,:self.K]
        self.X_df = df(self.X)
        # 去除缺失数据的行或列
        # axis=1/index 表示行   
        # axis=0/columns 表示列
        self.X_complete = self.X_df.dropna(axis=1)
        self.X_com = m(self.X_complete.values)
        # 得到丢失的数据
        self.X_missing = self.X_df.drop(self.X_complete.columns, axis=1)
        self.X_miss = m(self.X_missing.values)
        self.trajectory_dimentions = self.X_df.shape
        self.complete_dimensions = self.X_complete.shape
        self.missing_dimensions = self.X_missing.shape
        # 判断是否缺失数据，no_missing 是 boolean
        self.no_missing = self.missing_dimensions[1]==0
            
        if verbose:
            msg1 = 'Embedding dimension\t:  {}\nTrajectory dimensions\t: {}'
            msg2 = 'Complete dimension\t: {}\nMissing dimension     \t: {}'
            msg1 = msg1.format(self.embedding_dimension, self.trajectory_dimentions)
            msg2 = msg2.format(self.complete_dimensions, self.missing_dimensions)
            self._printer('EMBEDDING SUMMARY', msg1, msg2)
        
        if return_df:
            return self.X_df
    # 分解   
    def decompose(self, verbose=False):
        '''Perform the Singular Value Decomposition and identify the rank of the embedding subspace
        Characteristic of projection: the proportion of variance captured in the subspace'''
        X = self.X_com
        self.S = X*X.T
        self.U, self.s, self.V = linalg.svd(self.S)
        self.U, self.s, self.V = m(self.U), np.sqrt(self.s), m(self.V)
        # 返回秩
        self.d = np.linalg.matrix_rank(X)
        Vs, Xs, Ys, Zs = {}, {}, {}, {}
        for i in range(self.d):
            # Zs 没有什么用
            Zs[i] = self.s[i]*self.V[:,i]
            Vs[i] = X.T*(self.U[:,i]/self.s[i])  #Vi
            Ys[i] = self.s[i]*self.U[:,i]
            Xs[i] = Ys[i]*(m(Vs[i]).T) #Xi
        # 求贡献值
        self.Vs, self.Xs = Vs, Xs
        self.s_contributions = self.get_contributions(X, self.s, False)
        self.r = len(self.s_contributions[self.s_contributions>0])
        # sum求和
        self.r_characteristic = round((self.s[:self.r]**2).sum()/(self.s**2).sum(),4)
        # 提取出前r个U
        # 一个i对应 U所有行的第i列
        self.orthonormal_base = {i:self.U[:,i] for i in range(self.r)}

        if verbose:
            msg1 = 'Rank of trajectory\t\t: {}\nDimension of projection space\t: {}'
            msg1 = msg1.format(self.d, self.r)
            msg2 = 'Characteristic of projection\t: {}'.format(self.r_characteristic)
            self._printer('DECOMPOSITION SUMMARY', msg1, msg2)

    # 输出贡献值
    def view_s_contributions(self, adjust_scale=False, cumulative=False, return_df=False):
        '''View the contribution to variance of each singular value and its corresponding signal'''
        contribs = self.s_contributions.copy()
        contribs = contribs[contribs.Contribution!=0]
        if cumulative:
            # 贡献值累加
            contribs['Contribution'] = contribs.Contribution.cumsum()
        if adjust_scale:
            # 比例调整，更清楚的反映
            contribs = (1/contribs).max()*1.1-(1/contribs)
        ax = contribs.plot.bar(legend=False)
        ax.set_xlabel("Singular_i")
        ax.set_title('Non-zero{} contribution of Singular_i {}'.\
                     format(' cumulative' if cumulative else '', '(scaled)' if adjust_scale else ''))
        if adjust_scale:
            ax.axes.get_yaxis().set_visible(False)
        vals = ax.get_yticks()
        ax.set_yticklabels(['{:3.0f}%'.format(x*100) for x in vals])
        if show_flag:
           plt.show(ax)
        if return_df:
            return contribs
    

    def view_reconstruction(self, *hankel, names=None, return_df=True, plot=True, symmetric_plots=False):
        '''Visualise the reconstruction of the hankel matrix/matrices passed to *hankel'''
        hankel_mat = None
        for han in hankel:
            if isinstance(hankel_mat,m):
                hankel_mat = hankel_mat + han
            else: 
                hankel_mat = han.copy()
        self.hankel_full = self.diagonal_averaging(hankel_mat)
        RMSE = np.sqrt(sum((self.hankel_full.values-self.ts_v)**2)/len(self.ts_v))
        print("重构 与 原序列 的 RMSE = "+str(RMSE))


        title = 'Reconstruction of signal'
        if names or names==0: 
            title += ' associated with singular value{}: {}'
            title = title.format('' if len(str(names))==1 else 's', names)
        if plot:
            ax = self.hankel_full.plot(legend=False, title=title)
            if symmetric_plots:
                velocity = hankel_full.abs().max()[0]
                ax.set_ylim(bottom=-velocity, top=velocity)
            if show_flag:
                plt.show()
        #取出重构序列，用于预测
        print(self.hankel_full.shape[0])
        self.hankel_full = self.hankel_full.values
        if return_df:
            return self.hankel_full
    # 预测
    def _forecast_prep(self, singular_values=None):
        self.X_com_hat = np.zeros(self.complete_dimensions)
        self.verticality_coefficient = 0
        self.forecast_orthonormal_base = {}
        if singular_values:
            try:
                for i in singular_values:
                    self.forecast_orthonormal_base[i] = self.orthonormal_base[i]
            except:
                if singular_values==0:
                    self.forecast_orthonormal_base[0] = self.orthonormal_base[0]
                else:
                    raise('Please pass in a list/array of singular value indices to use for forecast')
        else:
            self.forecast_orthonormal_base = self.orthonormal_base

        self.R = np.zeros(self.forecast_orthonormal_base[0].shape)[:-1]
        for Pi in self.forecast_orthonormal_base.values():
            
            self.X_com_hat += Pi*Pi.T*self.X_com
            pi = np.ravel(Pi)[-1]
            self.verticality_coefficient += pi**2
            self.R += pi*Pi[:-1]
        self.R = m(self.R/(1-self.verticality_coefficient))
        #self.X_com_tilde 目前并未知道有什么用 
        #且self.X_com_tilde ！=  self.hankel_full
        self.X_com_tilde = self.diagonal_averaging(self.X_com_hat)


    def forecast_recurrent_using_reconstruct(self, steps_ahead=12, singular_values=None, plot=False, return_df=False, **plotargs):
        '''Forecast from last point of original time series up to steps_ahead using recurrent methodology
        This method also fills any missing data from the original time series.'''
        try:
            self.X_com_hat
        except(AttributeError):
            self._forecast_prep(singular_values)
        self.ts_forecast = np.array(self.ts_v[0])

        for i in range(1, self.ts_N+steps_ahead):
            try:
                if np.isnan(self.hankel_full[i]):
                    #缺值
                    x = self.R.T*m(self.ts_forecast[max(0,i-self.R.shape[0]): i]).T
                    self.ts_forecast = np.append(self.ts_forecast,x[0])
                else:
                    self.ts_forecast = np.append(self.ts_forecast,self.hankel_full[i])
            except(IndexError):
                #往后预测steps_ahead个值
                #其前d个数据的线性组合来计算
                #d = self.R.shape[0]
                x = self.R.T*m(self.ts_forecast[i-self.R.shape[0]: i]).T
                self.ts_forecast = np.append(self.ts_forecast, x[0])
       
        sum = 0
        components = [k for k in range(steps_ahead)]
        for k in components:
            a = self.ts_forecast[self.ts_N+k]
            b = self.org[self.ts_N:].values[k]
            sum += (a-b)**2 
        RMSE = np.sqrt(sum / steps_ahead)


        #print(self.ts_forecast[self.ts_N:])
        #print(type(self.ts_forecast[self.ts_N:]))
        #print(np.array(self.org[self.ts_N:].values))
        #print(type(self.org[self.ts_N:].values))
        #print(self.ts_forecast[self.ts_N:] - self.org[self.ts_N:].values)

        #RMSE = np.sqrt(sum((self.ts_forecast[self.ts_N:]-self.org[self.ts_N:].values)**2)/6)

        #minus = (self.ts_forecast[self.ts_N:] - self.org[self.ts_N:].values)
        #kRMSE =  np.sqrt((minus[:]**2).sum()/steps_ahead)
        self._printer("RMSE = " + str(RMSE));
        self.forecast_N = i+1
        new_index = pd.date_range(start=self.ts.index.min(),periods=self.forecast_N, freq=self.freq)
        forecast_df = df(self.ts_forecast, columns=['Forecast'], index=new_index)
        forecast_df['Original'] = self.org

        if plot:
            forecast_df.plot(title='Forecasted vs. original time series', **plotargs)
            if show_flag:
               plt.show()
        if return_df:
            return forecast_df
        
    def forecast_recurrent(self, steps_ahead=12, singular_values=None, plot=False, return_df=False, **plotargs):
        '''Forecast from last point of original time series up to steps_ahead using recurrent methodology
        This method also fills any missing data from the original time series.'''
        try:
            self.X_com_hat
        except(AttributeError):
            self._forecast_prep(singular_values)
        self.ts_forecast = np.array(self.ts_v[0])
        for i in range(1, self.ts_N+steps_ahead):
            try:
                if np.isnan(self.ts_v[i]):
                    #缺值
                    x = self.R.T*m(self.ts_forecast[max(0,i-self.R.shape[0]): i]).T
                    self.ts_forecast = np.append(self.ts_forecast,x[0])
                else:
                    self.ts_forecast = np.append(self.ts_forecast,self.ts_v[i])
            except(IndexError):
                #其前d个数据的线性组合来计算
                #d = self.R.shape[0]
                x = self.R.T*m(self.ts_forecast[i-self.R.shape[0]: i]).T
                self.ts_forecast = np.append(self.ts_forecast, x[0])
        
        sum = 0
        components = [k for k in range(steps_ahead)]
        for k in components:
            a = self.ts_forecast[self.ts_N+k]
            b = self.org[self.ts_N:].values[k]
            sum += (a-b)**2 
        RMSE = np.sqrt(sum / steps_ahead)
        self._printer("RMSE = " + str(RMSE));

        self.forecast_N = i+1
        #new_index = pd.date_range(start=self.ts.index.min(),periods=self.forecast_N, freq=self.freq)
        #forecast_df = df(self.ts_forecast, columns=['Forecast'], index=new_index)
        forecast_df['Original'] = np.append(self.ts_v, [np.nan]*steps_ahead)
        forecast_df['Original'] = self.org

        if plot:
            forecast_df.plot(title='Forecasted vs. original time series', **plotargs)
            if show_flag:
               plt.show()
        if return_df:
            return forecast_df
            
if __name__=='__main__':
    from mySSA_to_fit import mySSA_to_fit
    from pandas import DataFrame as df
    import pandas as pd
    import numpy as np
    from matplotlib.pylab import rcParams

    #总数据个数
    total_index = 195
    #去除数据的个数
    date_index = 15
    #预测步数  在此不需要用到
    steps_ahead = 6

    #窗口取值
    #file_index= ["29","33","36","48"]
    embedding_dimension_range= [17,20,23,24,27,29,33,36,42,46,48,51,56,58,60,64,80,96]

    #embedding_dimension_range= [12,19,22,29]
    for k in embedding_dimension_range:
        #设置窗口
        embedding_dimension=k
        #读取原数据并进行裁剪
        ts = pd.read_csv('Corn.csv', parse_dates=True, index_col='Month')
        ts_ = ts.copy()[:total_index-date_index]
        ssa = mySSA_to_fit(ts_,steps_ahead)
    
    
        #嵌入分解
        ssa.embed(embedding_dimension = embedding_dimension,  verbose=True)
        ssa.decompose(True)
       
       
        # 重构
        components = [i for i in range(ssa.r)]
        result  = ssa.view_reconstruction(*[ssa.Xs[i] for i in components], names=components)
        writer_date = df({"Price":""},index = ts_.index)
        writer_date['Price'] = result
        writer_date.to_csv("Corn"+str(embedding_dimension)+".csv")


        # 原序列和重构序列对比
        #rcParams['figure.figsize'] = 11,4
        #streams = [i for i in range(ssa.r)]
        #reconstructed = ssa.view_reconstruction(*[ssa.Xs[i] for i in streams], 
        #         names=streams, return_df=True, plot=False)
    
        #ts_copy = ssa.ts.copy()
        #ts_copy['Reconstruction'] = reconstructed
        #RMSE = np.sqrt(sum((ts_copy['Price']-ts_copy['Reconstruction'])**2)/len(ts_copy['Reconstruction']))
        #ts_copy.plot(title="Original vs. Reconstructed of "+str(ssa.r)+ " Time Series"
        #+"\nRMSE : "+str(RMSE))
        #plt.show(ts_copy.all)




        # FORECASTING
        #ssa.forecast_recurrent(steps_ahead, plot=True)
        #ssa.forecast_recurrent_using_reconstruct(steps_ahead, plot=True)